'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 1_simple_chat_bot.py
* @Time: 2025/10/23
* @All Rights Reserve By Brtc
'''
from multiprocessing.connection import answer_challenge
from typing import Any

import dotenv
from openai import OpenAI
dotenv.load_dotenv()
client = OpenAI(base_url="https://api.ephone.chat/v1")
#1、max_token用于判断是都徐要生成新的摘要
#2、summary 用于存储摘要
#3、chat_histories用于存储历史会话
#4、get_num_tokens用于计算传入的文本token数量
#5、save_context 用于存储新的交流对话
#6、get_buffer_string 用于讲历史对话转换成字符串
#7、load_memory_variables用于加载记忆变量信息
#8、summary_text用于将旧的摘要和传入对话生成新的摘要

class ConversationSummaryBufferMemory:
    def __init__(self, summary:str = "", chat_histories:list=None, max_tokens:int = 300):
        self._summary = summary
        self._chat_histories = chat_histories
        self._max_tokens = max_tokens
        self._client=OpenAI(base_url="https://api.ephone.chat/v1")


    def summary_text(self, org_summary:str, new_line:str)->str:
        """摘要生成函数老摘要  +  最远记录  = 新摘要"""
        prompt = f"""  
        你是一个强大的聊天机器人，请根据用户提供的谈话内容，总结摘要，并将其添加到先前提供的摘要中，返回一个新的摘要，除了新摘要其他任何数据都不要生成，如果用户的对话信息里有一些关键的信息，比方说姓名、爱好、性别、重要事件等等，这些全部都要包括在生成的摘要中，摘要尽可能要还原用户的对话记录。

        请不要将<example>标签里的数据当成实际的数据，这里的数据只是一个示例数据，告诉你该如何生成新摘要。

        <example>
        当前摘要：人类会问人工智能对人工智能的看法，人工智能认为人工智能是一股向善的力量。

        新的对话：
        Human：为什么你认为人工智能是一股向善的力量？
        AI：因为人工智能会帮助人类充分发挥潜力。

        新摘要：人类会问人工智能对人工智能的看法，人工智能认为人工智能是一股向善的力量，因为它将帮助人类充分发挥潜力。
        </example>

        =====================以下的数据是实际需要处理的数据=====================

        当前摘要：{org_summary}

        新的对话：
        {new_line}
        请帮用户将上面的信息生成新摘要。
        """
        completion = self._client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role":"user", "content":prompt}],
        )
        return completion.choices[0].message.content



    def get_buffer_string(self)->str:
        """获取历史消息"""
        buffer: str = ""
        for chat in self._chat_histories:
            buffer += f"Human:{chat.get('human')}\n AI:{chat.get('ai')}\n\n"
        return buffer.strip()


    def load_memory_variables(self)->dict[str, Any]:
        """获取整个记忆变量，方便添加到提示词中"""
        buffer_string = self.get_buffer_string() # 获取历史信息 "human:XXX ai:XXX"
        return {
            "chat_histories":f"摘要:{self._summary}\n\n 历史信息:{buffer_string}\n"
        }


    @classmethod
    def get_num_tokens(cls, query:str)->int:
        """获取历史消息的token数量我们这边为了方便就用 len来代替， 实际业务中请合理调整"""
        return len(query)


    def save_context(self, human_query:str, his_content:str)->None:
        """对话记录保存逻辑"""
        #1、先将对话添加到历史对话记录中去
        self._chat_histories.append({"human":human_query, "ai":his_content})
        #2、获取历史对话记录的总的字符串
        buffer_string = self.get_buffer_string()
        #3、获取历史对话记录token数
        tokens = self.get_num_tokens(buffer_string)
        #4、判断历史记录token数量是否超过了限定值
        if tokens > self._max_tokens:
            #5、获取最远的那一条记录
            first_chat = self._chat_histories[0]
            print("==================================> 开始生成新的摘要！")
            #6、生成新的摘要  = 老的摘要 + 最远的记录
            self._summary = self.summary_text(org_summary=self._summary,
                                              new_line=f"Human:{first_chat.get('human')}\n AI:{first_chat.get('ai')}\n\n")
            print("==================================> 摘要生成完毕！")
            # 删掉最远的记录
            del self._chat_histories[0]



if __name__ == "__main__":
    #定义一个对话缓冲窗口记忆类(参数是 摘要 + 历史记录 + tokens数量)
    app_mem = ConversationSummaryBufferMemory("", [], 300)

    while True:
        human_input = input("你好我是博小睿客服助手,请输入您的问题>")# 人类输入的问题
        if human_input == "exit":
            print("goodbye!!")
            exit(0)

        mem_var = app_mem.load_memory_variables()# 从记忆类中获取 摘要  + 记忆 字符字符串 格式:"摘要：xxxx 记录：AI:XXX   Human:XXX"
        answer_prompt = (
            "你是一个强大的聊天机器人，请根据用户上下文来回答问题。\n\n"
            f"{mem_var.get('chat_histories')}\n" # 历史记录添加到 提示词中
            f"用户的提问是:{human_input}"
        )
        completion = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role":"system", "content":"你是OpenAi研发聊天机器人"},
                {"role":"user", "content":answer_prompt},
            ],
            stream = True
        )
        print("博小睿:", end = "", flush=True)
        ai_content = ""
        for chunk in completion:
            content = chunk.choices[0].delta.content

            if content  is None:
                break
            ai_content += content
            print(content, end = "", flush=True)

        app_mem.save_context(human_input, ai_content)# 存储
        print("\n---------------------------------------------\n")

